import time
import os
import matplotlib.pyplot as plt
import numpy as np

# 导入你的加密算法实现
from publickey.rsa_core import RSA
from publickey.elgamal_core import ElGamal
from publickey.ecc_core import get_curve_by_name, generate_keypair as ecc_gen_keys, ecc_encrypt_simplified, ecc_decrypt_simplified

# ==============================================================================
# 1. 通用的性能测试辅助函数 (已更新以支持ElGamal)
# ==============================================================================
def run_single_performance_test(algo_type, config, num_iterations=100, message_size_bytes=40):
    """
    对单个算法配置进行性能测试的辅助函数。
    """
    results = {}
    print(f"\n--- 测试 {algo_type} with config: {config} ---")
    messages = [os.urandom(message_size_bytes) for _ in range(num_iterations)]

    # --- 密钥生成 ---
    start_time = time.perf_counter()
    if algo_type == "RSA":
        instance = RSA(bits=config)
        priv_key, pub_key = instance.private_key, instance.public_key
    elif algo_type == "ECC":
        instance = get_curve_by_name(config)
        priv_key, pub_key = ecc_gen_keys(instance)
    elif algo_type == "ElGamal":
        instance = ElGamal(key_size_bits=config)
        priv_key, pub_key = instance.private_key, instance.public_key
    else:
        raise ValueError("不支持的算法类型")
    results['key_gen_time'] = time.perf_counter() - start_time
    print(f"密钥生成时间: {results['key_gen_time']:.6f} 秒")
    
    # --- 加密 ---
    start_time = time.perf_counter()
    if algo_type == "RSA":
        ciphertexts = [instance.encrypt(m) for m in messages]
    elif algo_type == "ECC":
        ciphertexts = [ecc_encrypt_simplified(pub_key, instance, m) for m in messages]
    elif algo_type == "ElGamal":
        messages_int = [ElGamal.string_to_int(m.decode('latin-1')) for m in messages]
        ciphertexts = [ElGamal.encrypt(pub_key, m) for m in messages_int]
    results['encrypt_time'] = (time.perf_counter() - start_time) / num_iterations
    print(f"加密平均时间: {results['encrypt_time']:.8f} 秒")

    # --- 解密 ---
    start_time = time.perf_counter()
    if algo_type == "RSA":
        for c in ciphertexts: instance.decrypt(c)
    elif algo_type == "ECC":
        for R, C in ciphertexts: ecc_decrypt_simplified(priv_key, instance, R, C)
    elif algo_type == "ElGamal":
        public_params = (pub_key[0], pub_key[1], pub_key[2])
        for c in ciphertexts: ElGamal.decrypt(priv_key, public_params, c)
    results['decrypt_time'] = (time.perf_counter() - start_time) / num_iterations
    print(f"解密平均时间: {results['decrypt_time']:.8f} 秒")

    return results

# ==============================================================================
# 2. 两两对比的实现函数
# ==============================================================================

def create_comparison_plot(title, labels, data1, label1, data2, label2, filename):
    """一个通用的绘图函数，用于生成对比图表"""
    x = np.arange(len(labels))
    width = 0.35
    
    fig, ax = plt.subplots(figsize=(12, 7))
    rects1 = ax.bar(x - width/2, data1, width, label=label1, color='salmon')
    rects2 = ax.bar(x + width/2, data2, width, label=label2, color='skyblue')

    ax.set_ylabel('时间 (秒, 对数刻度)')
    ax.set_title(title)
    ax.set_xticks(x)
    ax.set_xticklabels(labels)
    ax.set_yscale('log')
    ax.legend()
    
    fig.tight_layout()
    plt.savefig(filename)
    plt.show()

def compare_rsa_vs_ecc():
    """在128-bit等效安全级别下比较RSA和ECC的性能"""
    print("\n\n" + "#"*15, " 对比 1: RSA vs ECC (128-bit 等效安全级别) ", "#"*15)
    
    rsa_results = run_single_performance_test("RSA", 3072)
    ecc_results = run_single_performance_test("ECC", "secp256r1")
    
    labels = ['Key Gen Time', 'Avg Encrypt Time', 'Avg Decrypt Time']
    rsa_values = [rsa_results['key_gen_time'], rsa_results['encrypt_time'], rsa_results['decrypt_time']]
    ecc_values = [ecc_results['key_gen_time'], ecc_results['encrypt_time'], ecc_results['decrypt_time']]

    create_comparison_plot(
        'RSA-3072 vs ECC-secp256r1 (128-bit安全级别) 性能对比',
        labels, rsa_values, 'RSA-3072', ecc_values, 'ECC-secp256r1',
        'rsa_vs_ecc_equivalence.png'
    )

# def compare_rsa_vs_elgamal():
#     """在相同密钥长度下比较RSA和ElGamal的性能"""
#     print("\n\n" + "#"*15, " 对比 2: RSA vs ElGamal (同密钥长度) ", "#"*15)
    
#     for bits in [1024, 2048]:
#         rsa_results = run_single_performance_test("RSA", bits)
#         elgamal_results = run_single_performance_test("ElGamal", bits)
        
#         labels = ['Key Gen Time', 'Avg Encrypt Time', 'Avg Decrypt Time']
#         rsa_values = [rsa_results['key_gen_time'], rsa_results['encrypt_time'], rsa_results['decrypt_time']]
#         elgamal_values = [elgamal_results['key_gen_time'], elgamal_results['encrypt_time'], elgamal_results['decrypt_time']]

#         create_comparison_plot(
#             f'RSA-{bits} vs ElGamal-{bits} 性能对比',
#             labels, rsa_values, f'RSA-{bits}', elgamal_values, f'ElGamal-{bits}',
#             f'rsa_vs_elgamal_{bits}.png'
#         )


def compare_rsa_vs_elgamal():
    """在相同密钥长度下比较RSA和ElGamal的性能"""
    print("\n\n" + "#"*15, " 对比 2: RSA vs ElGamal (同密钥长度) ", "#"*15)
    
    bits = 3072
    if bits == 3072:
        print("\n警告: 即将开始 3072-bit 测试，此过程可能非常缓慢，请耐心等待...")
        
        rsa_results = run_single_performance_test("RSA", bits)
        elgamal_results = run_single_performance_test("ElGamal", bits)
        
        labels = ['Key Gen Time', 'Avg Encrypt Time', 'Avg Decrypt Time']
        rsa_values = [rsa_results['key_gen_time'], rsa_results['encrypt_time'], rsa_results['decrypt_time']]
        elgamal_values = [elgamal_results['key_gen_time'], elgamal_results['encrypt_time'], elgamal_results['decrypt_time']]

        create_comparison_plot(
            f'RSA-{bits} vs ElGamal-{bits} 性能对比',
            labels, rsa_values, f'RSA-{bits}', elgamal_values, f'ElGamal-{bits}',
            f'rsa_vs_elgamal_{bits}.png'
        )    


def compare_elgamal_vs_ecc():
    """在128-bit等效安全级别下比较ElGamal和ECC的性能"""
    print("\n\n" + "#"*15, " 对比 3: ElGamal vs ECC (128-bit 等效安全级别) ", "#"*15)
    print("注意: ElGamal-3072 密钥生成会非常缓慢，请耐心等待...")
    
    elgamal_results = run_single_performance_test("ElGamal", 3072)
    ecc_results = run_single_performance_test("ECC", "secp256r1")
    
    labels = ['Key Gen Time', 'Avg Encrypt Time', 'Avg Decrypt Time']
    elgamal_values = [elgamal_results['key_gen_time'], elgamal_results['encrypt_time'], elgamal_results['decrypt_time']]
    ecc_values = [ecc_results['key_gen_time'], ecc_results['encrypt_time'], ecc_results['decrypt_time']]

    create_comparison_plot(
        'ElGamal-3072 vs ECC-secp256r1 (128-bit安全级别) 性能对比',
        labels, elgamal_values, 'ElGamal-3072', ecc_values, 'ECC-secp256r1',
        'elgamal_vs_ecc_equivalence.png'
    )

# ==============================================================================
# 3. 主程序入口
# ==============================================================================
if __name__ == "__main__":
    # 设置中文字体以正确显示图表
    plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei'] 
    plt.rcParams['axes.unicode_minus'] = False 
    
    # 依次运行所有对比测试
    compare_rsa_vs_ecc()
    compare_rsa_vs_elgamal()
    compare_elgamal_vs_ecc()
    
    print("\n\n所有对比演示已完成！")